import json
import multiprocessing
import os
import random
import re
import sys
import time
import winreg
import zipfile
import psutil
import requests
from lxml import etree
from selenium import webdriver
from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
from datetime import datetime, timedelta
from utils.utilities import message_head
from utils.general_settings import max_webdriver, chromedriver_path, chromedriver_url, zip_file_name


class TokenAssistant:
    @staticmethod
    def get_chrome_version():
        """获取已安装的chrome浏览器版本
        """
        try:
            key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, r'Software\Google\Chrome\BLBeacon')
            version, types = winreg.QueryValueEx(key, 'version')
            return version.split('.')[0]
        except Exception as e:
            print(message_head(1) + "获取Chrome浏览器版本出现错误！")
            print(message_head(1), e)
            return None

    @staticmethod
    def clear_process(pid, name):
        # 遍历当前全部进程
        for proc in psutil.process_iter():
            # 如果名称和pid均匹配则杀死进程
            if name in proc.name() and proc.pid == pid:
                proc.kill()

    @staticmethod
    def get_target_url():
        """通过给定的字符串列表生成一个url用于获取token
        """
        try:
            target = ['china', 'russia', 'usa', 'england', 'germany', 'france',
                      'beijing', 'berlin', 'paris', 'london', 'moscow', 'washington']
            rand = int(random.random() * 500 + 100)
            since_date = (datetime.now() - timedelta(days=rand)).strftime("%Y-%m-%d")
            until_date = (datetime.now() - timedelta(days=rand - 1)).strftime("%Y-%m-%d")
            url = str.format('https://twitter.com/search?q={0}%20since%3A{1}%20until%3A{2}&src=typed_query&f=live',
                             target[rand % len(target)], since_date, until_date)
            return url
        except Exception as e:
            print(message_head(1) + "获取目标URL出现错误！")
            print(message_head(1), e)
            return None

    @staticmethod
    def get_chromedriver_version(chromedriver_path):
        """获取指定位置的chromedriver的版本
        """
        try:
            if os.path.exists(chromedriver_path + "/chromedriver.exe"):
                version = os.popen(chromedriver_path + "/chromedriver.exe" + ' --version').read()
                return version.split(' ')[1].split('.')[0]
            else:
                print(message_head(0) + "Chromedriver文件不存在！")
                return None
        except Exception as e:
            print(message_head(1) + "获取Chromedriver版本出现错误！")
            print(message_head(1), e)
            return None

    @staticmethod
    def get_server_chrome_versions(chromedriver_url):
        """获取指定网址上的chromedriver的版本列表并返回一个version的字符串list
        chromedriver_url : 指定的chromedriver下载地址
        """
        version_dict = {}
        rep = requests.get(chromedriver_url).text
        html = etree.HTML(rep)
        for a in html.xpath('//a'):
            for href in a.xpath('@href'):
                version = re.compile(r'\d+\.\d+.\d+\.\d+').findall(href)  # 提取版本号
                if len(version) > 0 and '?' in href:
                    version_dict[version[0]] = href  # 将所有版本存入列表
        return version_dict

    @staticmethod
    def unzip(unzip_path, file_name):
        """解压Chromedriver压缩包到指定目录
        unzip_path : 指定的解压地址
        """
        f = zipfile.ZipFile(unzip_path + "/" + file_name, 'r')
        for file in f.namelist():
            f.extract(file, unzip_path)
        # print(message_head(0) + "解压成功")

    @staticmethod
    def get_webdriver_process_count(process_name):
        current_pid_list = psutil.pids()
        count = 0
        for pid in current_pid_list:
            try:
                process = psutil.Process(pid)
            except Exception as e:
                if 'PID not found' not in str(e):
                    print(message_head(1) + "获取webdriver进程数出错")
                    print(message_head(1), e)
                continue
            # 如果此pid对应的进程创建时间与传入的参数相同，则此进程为本程序创建且存活，保存信息
            if process_name in process.name():
                count += 1
        return count


class DriverUpdater:
    def __init__(self, chromedriver_path, chromedriver_url, zip_file_name):
        self.chromedriver_path = chromedriver_path
        self.chromedriver_url = chromedriver_url
        self.zip_file_name = zip_file_name

    def _download_chromedriver(self, download_url):
        """根据一个下载地址下载chromedriver
        """
        file = requests.get(download_url)
        with open(self.chromedriver_path + "/" + self.zip_file_name, 'wb') as zip_file:  # 保存文件到指定目录
            zip_file.write(file.content)
        # print(message_head(0) + "下载成功")

    def update_chromedriver(self, chrome_version):
        """解压Chromedriver压缩包到指定目录
        chrome_version : 本机安装的chrome版本
        chromedriver_path : 指定的保存地址
        """
        # print(message_head(0) + "正在更新Chromedriver.................")
        version_dict = TokenAssistant.get_server_chrome_versions(self.chromedriver_url)
        tmp_url = ""
        version_number = ''
        for version in version_dict.keys():
            if version.startswith(str(chrome_version)):
                version_number = version
                tmp_url = version_dict[version]
                break

        if len(version_number) > 0:
            download_url = str.format('https://chromedriver.storage.googleapis.com/{0}/chromedriver_win32.zip', version_number)
        else:
            print(message_head(1) + "未在概览页面找到与chrome匹配的chromedriver，请在【" + self.chromedriver_url + "】核实。")
            return
        print(message_head(0) + "正在下载【" + self.zip_file_name + "】，下载地址：", download_url)
        if len(download_url) > 0:
            self._download_chromedriver(download_url)
            TokenAssistant.unzip(self.chromedriver_path, self.zip_file_name)
            os.remove(self.chromedriver_path + "/" + self.zip_file_name)
        else:
            print(message_head(1) + "未在下载页面找到与chrome匹配的chromedriver，请检查【" + tmp_url + "】及【" + download_url + "】")
            time.sleep(300)
            return
        # print(message_head(0) + "更新后的Chromedriver版本为：",
        #       TokenAssistant.get_chromedriver_version(self.chromedriver_path))


class TokenProducer:
    """生成获取推文的request中的header所必须的token参数，一组token包含一个csrf和guest
    """

    def __init__(self, chromedriver_path, chromedriver_url, zip_file_name):
        """chromedriver_path : chromedriver程序本地保存的地址
        chromedriver_url : 提供chromedriver下载地址的页面
        download_url：chromedriver的实际下载地址，需要在网页上搜索后确定
        """
        self.chromedriver_path = chromedriver_path
        self.chromedriver_url = chromedriver_url
        self.zip_file_name = zip_file_name

    def get_token(self):
        """使用selenium获取一组token并返回数据
        """
        self._check_environment()

        webdriver_name = 'chromedriver.exe'
        chrome_options = webdriver.ChromeOptions()
        chrome_options.add_argument('--headless')  # 无界面
        chrome_options.add_argument('--no-sandbox')  # 解决DevToolsActivePort文件不存在报错问题
        chrome_options.add_argument('--disable-gpu')  # 禁用GPU硬件加速。如果软件渲染器没有就位，则GPU进程将不会启动。
        chrome_options.add_argument('--disable-dev-shm-usage')
        chrome_options.add_experimental_option('w3c', False)
        caps = DesiredCapabilities.CHROME
        caps['loggingPrefs'] = {'performance': 'ALL'}
        webdriver_pid = -1

        while True:
            # 如果现有的webdriver数量超过一定数值，则等待，否则启动webdriver。防止系统资源占用过大
            if TokenAssistant.get_webdriver_process_count(webdriver_name) >= min(multiprocessing.cpu_count(), max_webdriver):
                time.sleep(10)
            else:
                try:
                    driver = webdriver.Chrome(executable_path=self.chromedriver_path + "/" + webdriver_name,
                                              desired_capabilities=caps, options=chrome_options)
                    webdriver_pid = driver.service.process.pid
                    break
                except Exception as e:
                    print(message_head(1) + "ERROR:driver = webdriver.Chrome()")
                    print(message_head(1), e)
                    TokenAssistant.clear_process(webdriver_pid, webdriver_name)
                    time.sleep(60)

        extract_error_count = 0
        while True:
            try:
                driver.delete_all_cookies()
                driver.get(TokenAssistant.get_target_url())
                time.sleep(15)
                # 可能会出现这个问题，需要刷新页面
                while "Something went wrong, but don’t fret" in driver.page_source:
                    driver.refresh()
                    time.sleep(30)
                    extract_error_count += 1
                    if extract_error_count > 10:
                        print(message_head(0) + "连续10次刷新后出错！！！等待5分钟......")
                        extract_error_count = 0
                        time.sleep(300)
                break
            except Exception as e:
                print(message_head(1) + "driver.get(url)")
                print(message_head(1), e)
                TokenAssistant.clear_process(webdriver_pid, webdriver_name)
                time.sleep(60)
        try:
            log = driver.get_log('performance')
            guest = set()
            pattern = r"x-guest-token.:.\d{1,}"
            matches = re.findall(pattern, str(log))
            for i in range(len(matches)):
                guest.add(matches[i])
            driver.quit()
            if len(guest) > 0:
                return str(guest.pop()).split('":"')[1]
            else:
                print(message_head(1) + "未获取到csrf和guest，请检查！")
                return None
        except Exception as e:
            print(message_head(1) + "提取token出错")
            print(message_head(1), e)
            TokenAssistant.clear_process(webdriver_pid, webdriver_name)
            time.sleep(60)
            return None

    def _check_environment(self):
        """获取token
        """
        while True:
            # 先检查chrome_version和chromedriver_version
            try:
                chrome_version = TokenAssistant.get_chrome_version()
                chromedriver_version = TokenAssistant.get_chromedriver_version(self.chromedriver_path)
                # print(message_head(0) + "chrome_version: ", chrome_version, "\t\tchromedriver_version: ", chromedriver_version)
            except Exception as e:
                print(message_head(1) + "获取Chrome和Chromedriver版本出现错误")
                print(message_head(1), e)
                time.sleep(300)
                continue
            # chrome_version为空，可能是没安装chrome
            try:
                if chrome_version is None:
                    print(message_head(0) + "未检测到Chrome版本号，请检查是否已安装Chrome浏览器")
                    time.sleep(300)
                    continue
            except Exception as e:
                print(message_head(1) + "检查Chrome和Chromedriver出现错误")
                print(message_head(1), e)
                time.sleep(300)
                continue
            # 没有chromedriver_version，则自动下载
            try:
                if not os.path.exists(self.chromedriver_path):
                    os.mkdir(self.chromedriver_path)
                if chromedriver_version is None:
                    print(message_head(0) + "未检测到Chromedriver版本号")
                    DriverUpdater(self.chromedriver_path, self.chromedriver_url, self.zip_file_name).update_chromedriver(chrome_version)
                    continue
            except Exception as e:
                print(message_head(1) + "下载Chromedriver出错")
                print(message_head(1), e)
                time.sleep(300)
                continue
            # 版本号不匹配，则更新chromedriver_version，匹配则完成检测，可以启动
            try:
                if str(chrome_version) != str(chromedriver_version):
                    print(message_head(0) + "【chrome】和【chromedriver】版本不匹配，正在更新Chromedriver！")
                    DriverUpdater(self.chromedriver_path, self.chromedriver_url, self.zip_file_name).update_chromedriver(chrome_version)
                    continue
                else:
                    break
            except Exception as e:
                print(message_head(1) + "更新Chromedriver出错")
                print(message_head(1), e)
                time.sleep(300)
